(defpackage :linearize-transformation
  (:use :cl)
  (:shadow :condition)
  (:nicknames :lt)
  (:export :linearize-transformation
           :value-list
           :multiplier-list
           :condition
           :find-top-box
           :find-bottom-box))

(in-package :linearize-transformation)

(defclass linearize-transformation (o:transformation)
  ((value-list :accessor value-list)
   (multiplier-list :accessor multiplier-list)
   (condition :initarg :condition :accessor condition))
  (:default-initargs
   :surgeryp t
   :reward 4))

(defun part-of-a-loop-p (root node)
  "If NODE is part of a loop, return 3 values T, the base, and top of
the loop. NIL otherwise."
  (when (typep node 'd:argument)
    (setf node (tree:parent node)))
  (labels ((find-base (top)
             (u:let-if (base (d:find-if-on-tree (lambda (b)
                                                  (every (u:curry tree:descendant-p (p) b)
                                                         (tree:parents top)))
                                                root :only-boxes t
                                                :from-end t))
               (when (or (eq node base)
                         (tree:descendant-p node base))
                 base))))
    (u:let*-if ((top (d:find-if-on-tree (lambda (b)
                                          (and (d:multiparent-box-p b)
                                               (u:neq b node)
                                               (not (and
                                                     (d:xbox-p b)
                                                     (d:quantity (tree:child b))))
                                               (find-base b)))
                                        node :only-boxes t))
                (base (find-base top)))
      (values t top base))))

(defstruct box-loop
  base boxes top)

(defun eprint (node)
  (etypecase node
    (d:box
     (mapcar 'd:quantity (tree:children node)))
    (d:argument
     `(,(d:quantity node)))))

(defun paths-to (from to)
  "Find some of the paths from FROM to TO."
  (labels ((up (current visited)
             (unless (find current visited)
               (if (eq current to)
                   `((,to))
                   (loop for i in (tree:children current)
                      for u = (up i visited)
                      when u append (mapcar (lambda (x) (cons current x)) u)))))
           (down (current visited)
             (append
              (loop for i in (tree:parents current)
                 unless (typep i 'd:invisible-node)
                 append (loop
                           for j in (tree:children i)
                           for u = (up j (cons current visited))
                           append (loop for x in u collect (append visited (list current i) x))))
              (loop for i in (tree:parents current)
                 append (down i (cons current visited))))))
    (down from (list from))))

(defun find-all-loops (root)
  "Return all the loops in the tree starting at ROOT."
  (labels ((loops-for (top)
              (let* ((parents (tree:parents top))
                     (pairs (loop for i on parents
                               append (loop for j in (cdr i)
                                         collect (list (car i) j)))))
                (loop for i in pairs
                   for a = (first i)
                   for b = (second i)
                   append (loop for j in (paths-to a b)
                             collect (make-box-loop :base (find-if-not (lambda (b)
                                                                         (find-if (u:curry tree:descendant-p b (ancestor)) j))
                                                                       j)
                                                    :top top :boxes j))))))
    (sort
     (loop for mp in (d:collect-on-tree (lambda (b)
                                          (and (u:neq b root)
                                               (d:multiparent-box-p b)
                                               (not (and (d:xbox-p b)
                                                         (d:quantity (tree:child b))))))
                                        root :only-boxes t)
        append (loops-for mp))
     '< :key (lambda (a) (length (box-loop-boxes a))))))

(defun minimal-loop-p (l)
  (let ((boxes (remove-if-not (lambda (b) (typep b 'd:box)) (box-loop-boxes l))))
    (every (lambda (b)
             (let ((ret (and
                         (every (lambda (c)
                                  (or (d:quantity c)
                                      (and (tree:children c)
                                           (not (set-difference (tree:children c) (cons (box-loop-top l) boxes))))))
                                (tree:children b))
                         (or (and (eq b (box-loop-base l))
                                  (not (d:quantity (d:box-left-arg b)))
                                  (not (d:quantity (d:box-right-arg b))))
                             (and (d:quantity (first (tree:children b)))
                                  (not (d:quantity (third (tree:children b)))))
                             (and (not (d:quantity (first (tree:children b))))
                                  (d:quantity (third (tree:children b))))))))
               (or (eq b (box-loop-top l))
                   ret)))
           boxes)))

(defmethod o:transformation-matches-p ((trans linearize-transformation) root ignore-selected-p)
  (let ((loops (find-all-loops root)))
    (loop for l in loops
         for boxes = (remove-if-not (lambda (b) (typep b 'd:box)) (box-loop-boxes l))
       when (and (or ignore-selected-p
                     (every 'd:selected boxes))
                 (minimal-loop-p l))
       return (values t (box-loop-base l) `((loop ,l)) boxes))))

(defun get-answer (l)
  (labels ((go-up-side (box)
             (cond
               ((eq box (box-loop-top l))
                (values :empty (box-loop-top l)))
               ((and (d:quantity (d:box-left-arg box))
                     (d:quantity (d:box-right-arg box)))
                (error "I don't know how to handle this!"))
               ((d:quantity (d:box-left-arg box))
                (if (find (d:quantity (d:box-operator-arg box)) '(+ -))
                    (multiple-value-bind (num x-num) (go-up-side (d:box-right box))
                      (values (list (d:box-operator-arg box)
                                    (d:box-left-arg box)
                                    num)
                              (list (d:box-operator-arg box)
                                    x-num)))
                    (multiple-value-bind (num x-num) (go-up-side (d:box-right box))
                      (values (list (d:box-operator-arg box)
                                    (d:box-left-arg box)
                                    num)
                              (list (d:box-operator-arg box)
                                    (d:box-left-arg box)
                                    x-num)))))
               ((d:quantity (d:box-right-arg box))
                (if (find (d:quantity (d:box-operator-arg box)) '(+ -))
                    (multiple-value-bind (num x-num) (go-up-side (d:box-left box))
                      (values (list (d:box-operator-arg box)
                                    num
                                    (d:box-right-arg box))
                              ;; Make sure 4*(x+1) appears as (* (+ x)
                              ;; 4) so that it's distinguishable from
                              ;; 4*x.
                              (if (eq x-num (box-loop-top l))
                                  (list '+ x-num)
                                  x-num)))
                    (multiple-value-bind (num x-num) (go-up-side (d:box-left box))
                      (values (list (d:box-operator-arg box)
                                    num
                                    (d:box-right-arg box))
                              (list (d:box-operator-arg box)
                                    x-num
                                    (d:box-right-arg box))))))
               (t
                (let (num-1 x-num-1 num-2 x-num-2)
                  (multiple-value-setq (num-1 x-num-1) (go-up-side (d:box-left box)))
                  (multiple-value-setq (num-2 x-num-2) (go-up-side (d:box-right box)))
                  (values (list (d:box-operator-arg box)
                                num-1
                                num-2)
                          (list (d:box-operator-arg box)
                                x-num-1
                                x-num-2)))))))
      (go-up-side (box-loop-base l))))

(defun format-ans (arithmetic)
  (labels ((meat (tree)
             (cond
               ((atom tree)
                tree)
               ((= (length tree) 2)
                (let ((ans (meat (second tree))))
                  (cond
                    ((find (first tree) '(+ *))
                     ans)
                    ((numberp ans)
                     (funcall (first tree) ans))
                    ((eq ans :empty)
                     0)
                    (t
                     (print tree)
                     (format nil "~a~a" (first tree) (cond
                                                       ((eq ans :x)
                                                         1)
                                                       ((or
                                                         (and (listp (second tree))
                                                              (stringp ans)
                                                              (find (first (second tree)) '(+ - * /)))
                                                         ;; better safe than sorry
                                                         (stringp (second tree)))
                                                        (format nil "(~a)" ans))
                                                       (t ans)))))))
               ((and (find :empty tree)
                     (find (first tree) '(- +)))
                (if (and (eq (first tree) '-)
                         (u:neql (second tree) :empty))
                    (meat (list* '+ (cdr (remove :empty tree))))
                    (meat (remove :empty tree))))
               ((and (find :x tree)
                     (eq (first tree) '*))
                (meat (remove :x tree)))
               ((and (find :empty tree)
                     (eq (first tree) '*))
                :empty)
               (t
                (let ((items (loop for i in (cdr tree)
                                for new = (meat i)
                                collect (cond
                                          ((eq new :x)
                                           "1")
                                          ((and (listp i)
                                                (stringp new)
                                                (find (first i) '(+ -))
                                                (find (first tree) '(* /)))
                                           (format nil "(~a)" new))
                                          (t
                                           new)))))
                  (if (equal items (cdr tree))
                      (format nil (format nil "~~{~~a~~^~a~~}" (first tree)) items)
                      (meat (cons (car tree) items))))))))
    (let ((ret (meat arithmetic)))
      (cond
        ((eq ret :empty)
         0)
        ((eq ret :x)
         1)
        (t ret)))))

(defun test-format-ans ()
  (labels ((h (q a)
             (print (string= (u:string2 (format-ans q)) a))))
    (h :empty "0")
    (h :x "1")
    (h '(+ (+ (* 0 :x)) 1) "0+1")
    (h '(- :EMPTY (- 7 :EMPTY)) "-7")
    (h '(- :X (- :X)) "1--1")
    (h '(- (+ (* 9 :EMPTY) 9) (* :EMPTY 3)) "9")
    (h '(- (* 8 :EMPTY) :EMPTY) "0")
    (h '(- (* 7 :EMPTY) (- (+ 0 (* -1 :EMPTY)) 4)) "-(0-4)")
    (h '(- (* 4 :EMPTY)) "0")
    'done
    ))

(defstruct hint
  loop text which nodes)

(defmethod o:generate-new-tree ((trans linearize-transformation) vars)
  (labels ((strip (answer)
             (typecase answer
               (d:argument (d:quantity answer))
               (d:box :x)
               (list (mapcar #'strip answer))
               (t answer)))
           (ans (tree)
             (cond
               ((eq tree :empty)
                0)
               ((eq tree :x)
                1)
               ((atom tree)
                tree)
               (t
                (apply (car tree) (mapcar #'ans (cdr tree))))))
           (hint (answer nums loop which)
             (make-hint
              :text (ecase (condition trans)
                      (:discovery
                       "Enter a number.")
                      (:instruction
                       (let ((fmt (format-ans (strip nums))))
                         (if (numberp fmt)
                             (format nil "enter ~d because the sum of the ~a is ~d."
                                     answer
                                     (ecase which
                                       (:unknown "X's")
                                       (:constant "numbers"))
                                     answer)
                             (ecase which
                                       (:unknown
                                        (format nil "enter ~d because the sum of the X's is ~d ()."
                                                answer
                                                answer))
                                       (:constant
                                        (format nil "enter ~d because the sum of the numbers is ~d ()."
                                                answer
                                                answer)))))))
              :loop loop
              :nodes nums
              :which which)))
    (let ((l (second (find 'loop vars :key 'first))))
      (multiple-value-bind (num x-num) (get-answer l)
        ;; store how we got the answer in the transformation. FIXME: Total hack!
        ;;(print `(newt ,num ,x-num))
        (let* ((va (ans (strip num)))
               (ma (ans (strip x-num)))
               (v [i nil :answer va :hint (hint va num l :constant)])
               (m [i nil :answer ma :input-location :above :hint (hint ma x-num l :unknown)]))
          (setf (value-list trans) (list v num)
                (multiplier-list trans) (list m x-num))
          [m v [a '+] [a nil [m m [a '*] [a nil (box-loop-top l)]]]])))))

(defun display-loops (diagram nth)
  (let* ((loops (find-all-loops (d:root diagram)))
        (l (nth nth loops)))
    (d:clear-selected (d:root diagram))    
    (when l
      (setf (d:selected (box-loop-top l)) t)
      (dolist (i (box-loop-boxes l))
        (setf (d:selected i) t)))
    (setf (gwee:view-needs-update-p diagram) t)))

(defun find-bottom-box (root boxes)
  (d:find-if-on-tree (lambda (b)
                       (and (find b boxes)
                            (null (d:quantity (d:box-left-arg b)))
                            (null (d:quantity (d:box-right-arg b)))))
                     root
                     :only-boxes t))

(defun find-top-box (root boxes)
  (d:find-if-on-tree (lambda (b)
                       (and (not (find b boxes))
                            (d:parent-boxes b)
                            (= (count-if (lambda (arg)
                                           (and (find (tree:parent arg) boxes)
                                                (null (d:quantity arg))))
                                         (tree:parents b))
                               2)))
                     root
                     :only-boxes t))
